#!/usr/bin/env python
"""
exact_diag_calc_prop

"""


from __future__ import division
import numpy as np
import scipy
import matplotlib.pyplot as plt
import pickle


def get_partitionf(e, T):
    Z = np.sum(np.exp(-e/T))
    return Z

def get_energy(e, T, Z):
    E = np.sum(e*np.exp(-e/T))
    E = E/Z
    return E

def get_heat_capacity(e, T, E, Z):
    E2 = np.sum( (e**2) * np.exp(-e/T))/Z
    C = (E2 - E**2)/(T**2)
    return C

def get_susp_mag_vals(eigvec, eigval, N):
    i = 0
    n = np.size(eigval)
    X = np.zeros(n)
    M = np.zeros(n)
    for v in eigvec:
        X[i] = 0
        M[i] = 0
        for j in range(0, np.size(v)):
            M[i] += ( v[j]**2 ) * ( np.sum( get_Sz(n - j - 1, N) )/N )
            X[i] += ( v[j]**2 ) * ( (np.sum(get_Sz(n - j - 1, N))/N )**2 )
        i += 1
    return M, X

def get_avg(X, e, T, Z):
    X_ave = np.sum( X * np.exp(-e/T))/Z
    return X_ave


def get_Sz(i,N):
    Sz = np.zeros(N)
    b_str = np.binary_repr(i, N)
    for k in range(N):
        Sz[N - 1 - k] = np.int(b_str[k])
    Sz = 2 * Sz - 1
    return Sz


def get_state(Sz):
    N = np.size(Sz)
    Sz = (Sz + 1)/2
    i = np.int( np.sum(Sz * (2**np.arange(N))) )
    return i


N = 4
B = 0.0
temp_range = np.arange(0.01, 8, 0.2)
L = int( N**(1/2))
#print L
#print eig

prefix = 'exact_diag_N_%d_B_%0.2f_eig'%(N,B)
fname = prefix + '_values.pick'
eig = pickle.load(open(fname))
fname = prefix + '_vectors.pick'
eigvec = pickle.load(open(fname))

prefix = 'L_%d_B_%0.2f'%(L,B)

E_temp = []
heat_cap_temp = []
X_temp = []
#print "starting"
fname = 'susp_' + prefix + '.pick'
M, X = get_susp_mag_vals(eigvec, eig, N)

pickle.dump(X,open(fname,"wb"))
#print "done"
#X = pickle.load(open(fname))


for T in temp_range:
    Z = get_partitionf(eig, T)
    E = get_energy(eig, T, Z)
    E_temp.append(E/N)
    C = get_heat_capacity(eig, T, E, Z)
    heat_cap_temp.append(C/N)
    X_ave = get_avg(X, eig, T, Z)
    M_ave = get_avg(M, eig, T, Z)
    X_temp.append(X_ave)
    print T, E/N, C/N, X_ave, X_ave/T**2


prefix = prefix + '.png'
plt.plot(temp_range,E_temp,'b-')
fname = 'mean_energy_' + prefix
plt.savefig(fname)
plt.clf()
plt.plot(temp_range,heat_cap_temp,'b-')
fname = 'heat_cap_'+ prefix
plt.savefig(fname)
plt.clf()
plt.plot(temp_range,X_temp,'b-')
fname = 'susp_'+ prefix
plt.savefig(fname)
plt.clf()

